#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, os, re
import csv
from xml.dom.minidom import getDOMImplementation

TOOLS = [
		{'name':'MeCab', 'ver':'0.994'},
		{'name':'CaboCha', 'ver':'0.64'},
		{'name':'UniDic', 'ver':'1.3.12'},
		{'name':'ChaPAS', 'ver':'0.51'},
		{'name':'Modality', 'ver':'0.01'}
		]


# Global Regex patterns
ReArgId = re.compile('ID="([0-9]+)"', re.I)
ReArg = re.compile('(GA|WO|NI)="([0-9]+)"', re.I)
RePredType = re.compile('type="([a-zA-Z]+)"', re.I)



'''
parse mecab with unidic-parsed line to XML
'''
def parse_mecab_unidic(line, mid):
	attrs = ['pos1', 'pos2', 'pos3', 'pos4', 'cType', 'cForm', 'lForm', 'lemma', 'orth', 'pron', 'kana', 'goshu', 'orthBase', 'pronBase', 'kanaBase', 'formBase', 'iType', 'iForm', 'iConType', 'fType', 'fForm', 'fConType', 'aType', 'aConType', 'aModType']

	imp = getDOMImplementation()
	doc = imp.createDocument(None, None, None)
	token_ele = doc.createElement('Token')

	token_ele.setAttribute('id', 't%d' % mid)

	l = line.split('\t')
	reader = csv.reader([l[1]])
	for row in reader:
		i = 0
		for val in row:
			token_ele.setAttribute(attrs[i], val.decode('utf-8'))
			i += 1
	
	token_ele.setAttribute('surf', l[0].decode('utf-8'))
	token_ele.setAttribute('orig', token_ele.getAttribute('orthBase'))
	token_ele.setAttribute('pos', token_ele.getAttribute('pos1'))

	if len(l) == 3:
		token_ele.setAttribute('ne', l[2].decode('utf-8'))

	return token_ele


'''
parse chasen with ipadic-parsed line to XML
'''
def parse_chasen_ipadic(line, mid):
	attrs = ['surf', 'pron', 'orig', 'pos1', 'pos2', 'pos3']

	imp = getDOMImplementation()
	doc = imp.createDocument(None, None, None)
	token_ele = doc.createElement('Token')

	token_ele.setAttribute('id', 't%d' % mid)

	l = line.split('\t')
	for i, attr in enumerate(attrs):
		token_ele.setAttribute(attr, l[i].decode('utf-8'))

	token_ele.setAttribute('pos', token_ele.getAttribute('pos1').split('-')[0])

	if len(l) == 3:
		token_ele.setAttribute('ne', l[2].decode('utf-8'))

	return token_ele


'''
parse mecab with ipadic-parsed line to XML
'''
def parse_mecab_ipadic(line, mid):
	attrs = ['pos', 'pos1', 'pos2', 'pos3', 'cForm', 'cType', 'orig', 'read', 'pron']

	imp = getDOMImplementation()
	doc = imp.createDocument(None, None, None)
	token_ele = doc.createElement('Token')

	token_ele.setAttribute('id', 't%d' % mid)

	l = line.split('\t')
	reader = csv.reader([l[1]])
	for row in reader:
		i = 0
		for val in row:
			if len(attrs) > i:
				token_ele.setAttribute(attrs[i], val.decode('utf-8'))
			i += 1
	
	token_ele.setAttribute('surf', l[0].decode('utf-8'))

	if len(l) == 3:
		token_ele.setAttribute('ne', l[2].decode('utf-8'))

	return token_ele



'''
parse modality-analyzed lines to XML dict with token id as key
'''
def parse_modality(lines):
	modalities = {}

	for line in lines:
		if line.startswith('#EVENT'):
			l = line.split('\t')
			mids = [int(mid) for mid in l[1].split(',')]
			imp = getDOMImplementation()
			doc = imp.createDocument(None, None, None)
			mod_ele = doc.createElement('Modality')
			mod_ele.setAttribute('source', l[2])
			mod_ele.setAttribute('tense', l[3])
			mod_ele.setAttribute('assumptional', l[4])
			mod_ele.setAttribute('type', l[5])
			mod_ele.setAttribute('authenticity', l[6])
			mod_ele.setAttribute('sentiment', l[7])
			mod_ele.setAttribute('focus', l[8])
			for mid in mids:
				modalities[mid] = mod_ele

	return modalities


'''
parse predicate-argument analyzed results with SynCha format in token line
'''
def parse_pas(lines):
	mid = 0
	arg2mid = {}
	pas = {}
	for line in lines:
		if line.startswith('#') or line.startswith('EOS') or line == '' or line.startswith('* '):
			pass
		else:
			l = line.split('\t')
			if len(l) >= 8:
				for pas_info in l[7].split(' '):
					if pas_info.startswith('ID='):
						re_arg_id = ReArgId.match(pas_info)
						arg_id = re_arg_id.group(1)
						arg2mid[arg_id] = mid
					elif pas_info.startswith('GA=') or pas_info.startswith('WO=') or pas_info.startswith('NI='):
						re_arg = ReArg.match(pas_info)
						arg_type = re_arg.group(1)
						arg_id = re_arg.group(2)
						if not pas.has_key(mid):
							pas[mid] = {}
						pas[mid][arg_type.lower()] = arg_id
					elif pas_info.startswith('type='):
						re_pred_type = RePredType.match(pas_info)
						pred_type = re_pred_type.group(1)
						if not pas.has_key(mid):
							pas[mid] = {}
						pas[mid]['type'] = pred_type
			mid += 1

	pas_xmls = {}
	for mid, pas_info in pas.iteritems():
		imp = getDOMImplementation()
		doc = imp.createDocument(None, None, None)
		pred_ele = doc.createElement('Predicate')
		for key, val in pas_info.iteritems():
			if key == 'type':
				pred_ele.setAttribute(key, val)
			else:
				pred_ele.setAttribute(key, 't%d' % arg2mid[val])
		pas_xmls[mid] = pred_ele
	
	return pas_xmls


def make_xml(lines, id='null', role=None, morph_dic='ipadic', morph_tool='mecab'):
	str_xml = '<?xml version="1.0" encoding="UTF-8" standalone="no" ?>'
	str_doctype = '<!DOCTYPE dataset PUBLIC "-//Tohoku University//DTD statistics for sentence in RTE//EN" "http://www.cl.ecei.tohoku.ac.jp/rite2/dtd/sent.dtd">'

	imp = getDOMImplementation()
	doctype = imp.createDocumentType(
			qualifiedName = 'Sentence',
			publicId = 'http://www.cl.ecei.tohoku.ac.jp/rite2/dtd/sent.dtd',
			systemId = ''
			)
	doc = imp.createDocument(None, 'Sentence', doctype)

	root = doc.documentElement
	root.setAttribute('id', id)
	if role:
		root.setAttribute('role', role)

	for tool in TOOLS:
		ann_ele = doc.createElement('Annotation')
		ann_ele.setAttribute('tool', tool['name'])
		ann_ele.setAttribute('ver', tool['ver'])
		root.appendChild(ann_ele)

	# parse modelity information (Matsuyoshi's format)
	modalities = parse_modality(lines)

	# parse predicate-argument structure information (SynCha format)
	pas = parse_pas(lines)

	# parse chunk and token information
	mid = 0
	chunk_ele = None
	sent = ''
	for line in lines:
		if line.startswith('#') | line.startswith('EOS') | (line == ''):
			pass
		elif line.startswith('*'):
			if chunk_ele:
				root.appendChild(chunk_ele)
			chunk_ele = doc.createElement('Chunk')

			l = line.split(' ')
			chunk_ele.setAttribute('id', 'c%s' % l[1])
			head = l[2][0:len(l[2])-1]
			chunk_ele.setAttribute('head', 'c%s' % head)
			type = l[2][len(l[2])-1:len(l[2])]
			chunk_ele.setAttribute('type', type)
			chunk_ele.setAttribute('score', l[4])
			chunk_ele.setAttribute('func', l[3])
		else:
			token_ele = None
			if morph_tool == 'mecab':
				if morph_dic == 'ipadic':
					token_ele = parse_mecab_ipadic(line, mid)
				elif morph_dic == 'unidic':
					token_ele = parse_mecab_unidic(line, mid)
			elif morph_tool == 'chasen':
				if morph_dic == 'ipadic':
					token_ele = parse_chasen_ipadic(line, mid)
				elif morph_dic == 'unidic':
					sys.stderr.write('sorry, setting of morph tool: %s and morph dic: %s is comming soon\n' % (morph_tool, morph_dic))
					exit()
					token_ele = parse_chasen_unidic(line, mid)

			if modalities.has_key(mid):
				token_ele.appendChild(modalities[mid])

			if pas.has_key(mid):
				token_ele.appendChild(pas[mid])

			chunk_ele.appendChild(token_ele)
			mid += 1

			sent += line.split('\t')[0]

	# add last chunk
	root.appendChild(chunk_ele)

	root.setAttribute('text', sent)

	xml = doc.toprettyxml('  ', '\n', encoding='utf-8')

	p = re.compile('<\?xml[^>]*\?>')
	xml = p.sub(str_xml, xml)

	p = re.compile('<!DOCTYPE[^>]*>', re.M)
	xml = p.sub(str_doctype, xml)

	return xml


if __name__ == '__main__':
	import optparse

	MorphDics = ['ipadic', 'unidic']
	MorphTools = ['mecab', 'chasen']

	optparser = optparse.OptionParser()
	optparser.add_option('-i', '--id', action='store', type='string', dest='id', default='null', help='ID of sentence [optional] default:null')
	optparser.add_option('-r', '--role', action='store', type='string', dest='role', default='none', help='role of sentence [optional] default:none')
	optparser.add_option('-m', '--ma', action='store', type='string', default='chasen', dest='morph_tool', help='morphologic analysis tool name [optional] (%s) default:chasen' % ','.join(MorphTools))
	optparser.add_option('-d', '--dicname', action='store', type='string', default='ipadic', dest='morph_dic', help='dictionary name of morphologic analysis [optional] (%s) default:ipadic' % ','.join(MorphDics))

	(opts, args) = optparser.parse_args(sys.argv)

	lines = []
	for line in sys.stdin:
		lines.append(line.rstrip())

	xml = make_xml(lines, id=opts.id, role=opts.role, morph_dic=opts.morph_dic, morph_tool=opts.morph_tool)

	print xml

